# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License

"""ParquetTableEmitter module."""
from dataclasses import dataclass
from typing import Optional,Tuple,Dict,Union
import logging
import traceback
from pathlib import Path
import pandas as pd
from pyarrow.lib import ArrowInvalid, ArrowTypeError

from graphrag.index.storage import PipelineStorage
from graphrag.index.typing import ErrorHandlerFn
from .table_emitter import TableEmitter
log = logging.getLogger(__name__)
class ParquetTableEmitter(TableEmitter):
    """ParquetTableEmitter class."""

    _storage: PipelineStorage
    _on_error: ErrorHandlerFn
    _last_filename: Optional[str] = None  
    _last_filepath: Optional[Path] = None 

    def __init__(
        self,
        storage: PipelineStorage,
        on_error: ErrorHandlerFn,
    ):
        """Create a new Parquet Table Emitter."""
        self._storage = storage
        self._on_error = on_error
    async def emit(self, name: str, data: pd.DataFrame) -> None:
        """Emit a dataframe to storage."""
        filename = f"{name}.parquet"
        log.info("emitting parquet table %s", filename)
        root_dir = self._storage.get_root_dir()  # 直接调用方法
        filepath = Path(root_dir) / filename
        #print(filepath)#test815/output/20240815-113826/artifacts/create_base_text_units.parquet
        try:
            await self._storage.set(filename, data.to_parquet())
            self._last_filename = filename  
            self._last_filepath = filepath 
        except ArrowTypeError as e:
            log.exception("Error while emitting parquet table")
            self._on_error(
                e,
                traceback.format_exc(),
                None,
            )
        except ArrowInvalid as e:
            log.exception("Error while emitting parquet table")
            self._on_error(
                e,
                traceback.format_exc(),
                None,
            )
    def get_last_emitted_info(self) -> Optional[Dict[str, Optional[Union[str, Path]]]]:
        """Get the last emitted filename and filepath."""
        """字典的键是str，值可以是None，字符串或者path对象"""
        return {"filename": self._last_filename, "full_path": self._last_filepath}